package com.bruce.ai.alibaba.workflow.config;

import java.util.HashMap;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import com.alibaba.cloud.ai.graph.KeyStrategy;
import com.alibaba.cloud.ai.graph.KeyStrategyFactory;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.state.strategy.AppendStrategy;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;
import com.bruce.ai.alibaba.workflow.node.LongMemoryNode;
import com.bruce.ai.alibaba.workflow.node.MemoryNode;
import static com.alibaba.cloud.ai.graph.action.AsyncNodeAction.node_async;

@Configuration
public class MemoryWorkflowConfig {

    /**
     * 记忆工作流
     * 
     * @return
     */
    @Bean
    public StateGraph memoryWorkflow(ChatClient.Builder chatClientBuilder) throws Exception {
        KeyStrategyFactory keyStrategyFactory = () -> {
            HashMap<String, KeyStrategy> strategies = new HashMap<>();
            strategies.put("messages", new AppendStrategy());
            return strategies;
        };

        // 定义工作流
        StateGraph graph = new StateGraph(keyStrategyFactory);

        // 定义节点
        graph.addNode("chat", node_async(new MemoryNode(chatClientBuilder)));

        // 定义边
        graph.addEdge(StateGraph.START, "chat");
        graph.addEdge("chat", StateGraph.END);

        return graph;
    }

    @Bean
    public StateGraph longMemoryWorkflow(ChatClient.Builder chatClientBuilder) throws Exception {
        KeyStrategyFactory keyStrategyFactory = () -> {
            HashMap<String, KeyStrategy> strategies = new HashMap<>();
            strategies.put("userId", new ReplaceStrategy());
            strategies.put("userProfile", new ReplaceStrategy());
            return strategies;
        };

        // 定义工作流
        StateGraph graph = new StateGraph(keyStrategyFactory);

        // 定义节点
        graph.addNode("load_profile", com.alibaba.cloud.ai.graph.action.AsyncNodeActionWithConfig.node_async(new LongMemoryNode()));

        // 定义边
        graph.addEdge(StateGraph.START, "load_profile");
        graph.addEdge("load_profile", StateGraph.END);
        return graph;
    }
}
